# 2021-10-29
# 基于P_HASH的PRF
# 伪随机函数生成方法
# bo.li.rd@watchdata.com


import hmac
import math


class PRF:
    #每一轮输出的字节数
    LOOP_SIZE = 32

    __label = b'my_PRF_with_P_HASH'
    __seed = b''
    __secret = b''

    def __init__(self, seed:bytes, secret:bytes) -> bytes:
        self.__seed = self.__label + seed
        self.__secret = secret

    def __next_A(self, last_A):
        if not last_A:
            return self.__seed
        else:
            return hmac.new(self.__secret, last_A, digestmod='SHA256').digest()

    def __P_Hash(self, loops):
        out = b''
        last_A = b''
        for i in range(loops):
            last_A = self.__next_A(last_A)
            new_seed = last_A + self.__seed
            temp = hmac.new(self.__secret, new_seed, digestmod='SHA256').digest()
            out += temp
        return out

    def get(self, len):
        loops = math.ceil(len / self.LOOP_SIZE)
        out = self.__P_Hash(loops)
        return out[:len]


if __name__ == '__main__':

    p = PRF(b'abc', b'123')
    o = p.get(1)
    print(o.hex())
    o = p.get(32)
    print(o.hex())
    o = p.get(66)
    print(o.hex())